"""
各种常用的函数
运算函数，激活函数
"""

import numpy as np

import dezero.core
from dezero import utils
from dezero.core import *
from dezero import cuda
from dezero.functions_conv import *

#==============================一般函数
#具体实现类 实现平方函数
class Square(Function):
    def forward(self,x):
        y = x**2
        return y
    def backward(self,gy):
        x, = self.inputs
        gx = 2 * x * gy #x**2的微分式子
        return gx
def square(x):
    return Square()(x)
    
#具体实现类 实现exp函数
class Exp(Function):
    def forward(self,x):
        xp = cuda.get_array_module(x)
        y = xp.exp(x)
        return y
    def backward(self,gy):
        x, = self.inputs
        xp = cuda.get_array_module(x)
        gx = xp.exp(x) * gy
        return gx
def exp(x):
    return Exp()(x)

class Log(Function):
    def forward(self, x):
        xp = cuda.get_array_module(x)
        y = xp.log(x)
        return y

    def backward(self, gy):
        x, = self.inputs
        gx = gy / x
        return gx

def log(x):
    return Log()(x)
    
#==========================三角函数
class Sin(Function):
    def forward(self,x):
        xp = cuda.get_array_module(x)
        y = xp.sin(x)
        return y
    def backward(self,gy):
        x, = self.inputs
        gx = gy * cos(x)
        return gx
def sin(x):
    return Sin()(x)

class Cos(Function):
    def forward(self,x):
        xp = cuda.get_array_module(x)
        y = xp.cos(x)
        return y
    def backward(self,gy):
        x, = self.inputs
        gx = gy * -sin(x)
        return gx
def cos(x):
    return Cos()(x)

class Tanh(Function):
    def forward(self,x):
        xp = cuda.get_array_module(x)
        y = xp.tanh(x)
        return y
    def backward(self,gy):
        #tanh求导为 y' = 1 - y^2
        y = self.outputs[0]()
        gx = gy * (1 - y*y)
        return gx
def tanh(x):
    return Tanh()(x)


#====================张量处理(类np)
class Reshape(Function):
    def __init__(self,shape):
        self.shape = shape
    def forward(self,x):
        self.x_shape = x.shape
        y = x.reshape(self.shape)
        return y
    def backward(self,gy):
        return reshape(gy, self.x_shape)
def reshape(x,shape):
    if x.shape == shape:
        return as_variable(x)
    return Reshape(shape)(x)

class Transpose(Function):
    def __init__(self, axes=None):
        self.axes = axes

    def forward(self, x):
        y = x.transpose(self.axes)
        return y

    def backward(self, gy):
        if self.axes is None:
            return transpose(gy)

        axes_len = len(self.axes)
        inv_axes = tuple(np.argsort([ax % axes_len for ax in self.axes]))
        return transpose(gy, inv_axes)


def transpose(x, axes=None):
    return Transpose(axes)(x)

class Sum(Function):
    def __init__(self,axis,keepdims):
        self.axis = axis
        self.keepdims = keepdims
    def forward(self,x):
        self.x_shape = x.shape
        y = x.sum(axis=self.axis, keepdims=self.keepdims)
        return y
    def backward(self,gy):
        gy = utils.reshape_sum_backward(gy,self.x_shape,self.axis,self.keepdims)
        gx = broadcast_to(gy, self.x_shape)
        return gx
def sum(x,axis=None,keepdims=False):
    return Sum(axis,keepdims)(x)

class BroadcastTo(Function):
    def __init__(self,shape):
        self.shape = shape
    def forward(self,x):
        self.x_shape = x.shape
        xp = cuda.get_array_module(x)
        y = xp.broadcast_to(x, self.shape)
        return y
    def backward(self,gy):
        gx = sum_to(gy, self.x_shape)
        return gx
def broadcast_to(x,shape):
    if x.shape == shape:
        return as_variable(x)
    return BroadcastTo(shape)(x)

class SumTo(Function):
    def __init__(self,shape):
        self.shape = shape
    def forward(self,x):
        self.x_shape = x.shape
        y = utils.sum_to(x, self.shape)
        return y
    def backward(self,gy):
        gx = broadcast_to(gy,self.x_shape)
        return gx
def sum_to(x,shape):
    if x.shape == shape:
        return as_variable(x)
    return SumTo(shape)(x)


class MatMul(Function):
    def forward(self,x,W):
        y = x.dot(W)
        return y
    def backward(self,gy):
        x, W = self.inputs
        gx = matmul(gy, W.T)
        gW = matmul(x.T,gy)
        return gx,gW
def matmul(x,W):
    return MatMul()(x,W)

class GetItem(Function):
    def __init__(self, slices):
        self.slices = slices

    def forward(self, x):
        y = x[self.slices]
        return y

    def backward(self, gy):
        x, = self.inputs
        f = GetItemGrad(self.slices, x.shape)
        return f(gy)

class GetItemGrad(Function):
    def __init__(self, slices, in_shape):
        self.slices = slices
        self.in_shape = in_shape

    def forward(self, gy):
        xp = cuda.get_array_module(gy)
        gx = xp.zeros(self.in_shape, dtype=gy.dtype)
        xp.add.at(gx, self.slices, gy)
        return gx

    def backward(self, ggx):
        return get_item(ggx, self.slices)

def get_item(x, slices):
    f = GetItem(slices)
    return f(x)



# =============================================================================
# max / min / clip
# =============================================================================
class Max(Function):
    def __init__(self, axis=None, keepdims=False):
        self.axis = axis
        self.keepdims = keepdims

    def forward(self, x):
        y = x.max(axis=self.axis, keepdims=self.keepdims)
        return y

    def backward(self, gy):
        x = self.inputs[0]
        y = self.outputs[0]()  # weakref

        shape = utils.max_backward_shape(x, self.axis)
        gy = reshape(gy, shape)
        y = reshape(y, shape)
        cond = (x.data == y.data)
        gy = broadcast_to(gy, cond.shape)
        return gy * cond


class Min(Max):
    def forward(self, x):
        y = x.min(axis=self.axis, keepdims=self.keepdims)
        return y


def max(x, axis=None, keepdims=False):
    return Max(axis, keepdims)(x)


def min(x, axis=None, keepdims=False):
    return Min(axis, keepdims)(x)

class Clip(Function):
    """numpy.clip 裁剪"""
    def __init__(self, x_min, x_max):
        self.x_min = x_min
        self.x_max = x_max

    def forward(self, x):
        xp = cuda.get_array_module(x)
        y = xp.clip(x, self.x_min, self.x_max)
        return y

    def backward(self, gy):
        x, = self.inputs
        mask = (x.data >= self.x_min) * (x.data <= self.x_max)
        gx = gy * mask
        return gx

def clip(x, x_min, x_max):
    return Clip(x_min, x_max)(x)

#======================误差函数、感知器
class MeanSquaredError(Function):
    """线性回归MSE函数"""
    def forward(self,x0,x1):
        diff = x0 - x1
        y =  (diff ** 2).sum() / len(diff)
        return y
    def backward(self,gy):
        x0,x1 = self.inputs
        diff = x0 - x1
        gx0 = gy * diff * ( 2.0 / len(diff) )
        gx1 = -gx0
        return gx0,gx1
def mean_squared_error(x0, x1):
    return MeanSquaredError()(x0,x1)


def linear_simple(x,W,b=None):
    t = matmul(x,W)
    if b is None:
        return t
    y = t + b
    t.data = None #删除t的数据 节省内存
    return y

class Linear(Function):
    """线性变换 y=matmul(x,W)+b"""
    def forward(self,x,W,b):
        y = x.dot(W)
        if b is not None:
            y += b
        return y
    def backward(self,gy):
        x,W,b = self.inputs
        gx = matmul(gy,W.T)
        gW = matmul(x.T,gy)
        gb = None if b.data is None else sum_to(gy,b.shape) #正向传播是广播，反向是sum
        return gx,gW,gb
def linear(x,W,b=None):
    return Linear()(x,W,b)


#===============================
# 激活函数相关 Sigmoid ReLU Softmax
#================================

# Sigmoid
def sigmoid_simple(x):
    x = as_variable(x)
    y = 1 / ( 1 + exp(-x) )
    return y

class Sigmoid(Function):
    def forward(self, x):
        # y = 1 / (1 + np.exp(-x))
        xp = cuda.get_array_module(x)
        y = xp.tanh(x * 0.5) * 0.5 + 0.5  # Better implementation
        return y
    def backward(self, gy):
        y = self.outputs[0]()
        gx = gy * y * (1 - y)
        return gx
def sigmoid(x):
    return Sigmoid()(x)

#ReLU
class ReLU(Function):
    """ReLu函数在输入大于0时按原样输出，输入小于等于0时输出0"""
    def forward(self,x):
        xp = cuda.get_array_module(x)
        y = xp.maximum(x,0.0) #np.maximum(X, Y) 用于逐元素比较两个array的大小
        return y
    def backward(self,gy):
        x, = self.inputs
        mask = x.data > 0
        gx = gy * mask
        return gx
def relu(x):
    return ReLU()(x)


# Softmax
def softmax_simple(x, axis=1):
    x = as_variable(x)
    y = exp(x)
    sum_y = sum(y, axis=axis, keepdims=True)
    return y / sum_y

class Softmax(Function):
    def __init__(self, axis=1):
        self.axis = axis

    def forward(self, x):
        xp = cuda.get_array_module(x)
        y = x - x.max(axis=self.axis, keepdims=True)
        y = xp.exp(y)
        y /= y.sum(axis=self.axis, keepdims=True)
        return y

    def backward(self, gy):
        y = self.outputs[0]()
        gx = y * gy
        sumdx = gx.sum(axis=self.axis, keepdims=True)
        gx -= y * sumdx
        return gx

def softmax(x, axis=1):
    return Softmax(axis)(x)

def softmax_cross_entropy_simple(x,t):
    """
    描述：softmax + 交叉熵误差
    参数：x——输入数据，t——训练数据记录(内容是正确类别的布尔索引)
    """
    x,t = as_variable(x),as_variable(t)
    N = x.shape[0]

    p = softmax(x)
    p = clip(p, 1e-15, 1.0) #为了防止log(0)，将p限制为大于1e-15的值
    log_p = log(p)
    xp = cuda.get_array_module(x)
    tlog_p = log_p[xp.arange(N), t.data]
    y = -1 * sum(tlog_p) / N
    return y

class SoftmaxCrossEntropy(Function):
    def forward(self, x, t):
        xp = cuda.get_array_module(x)
        N = x.shape[0]
        log_z = utils.logsumexp(x, axis=1)
        log_p = x - log_z
        log_p = log_p[xp.arange(N), t.ravel()]
        y = -log_p.sum() / xp.float32(N)
        return y

    def backward(self, gy):
        x, t = self.inputs
        N, CLS_NUM = x.shape

        gy *= 1/N
        y = softmax(x)
        # convert to one-hot
        # xp = cuda.get_array_module(t.data)
        xp = cuda.get_array_module(x)
        t_onehot = xp.eye(CLS_NUM, dtype=t.dtype)[t.data]
        y = (y - t_onehot) * gy
        return y


def softmax_cross_entropy(x, t):
    return SoftmaxCrossEntropy()(x, t)


#=========================
# 评估精确度
def accuracy(y,t):
    """
    评估识别精度
    参数：
        y：神经网络的预测结果
        t：正确答案的数据
    原理：取出预测结果的最大值的索引，得到预测结果pred，
    然后将pred与数据t进行比较结果是True/False张量，
    计算张量为True的数据所占的比例(求平均值)得到的值就相当于正确率
    """
    y,t = as_variable(y), as_variable(t)
    pred = y.data.argmax(axis=1).reshape(t.shape) #找出神经网络预测结果的最大索引
    result = (pred == t.data) #result 为 True/False张量
    acc = result.mean()
    return Variable(as_array(acc))


#======================================
# 抑制过拟合 Dropout
#======================================

def dropout(x, dropout_ratio=0.5):
    """使用 反向Dropout"""
    xp = cuda.get_array_module(x)
    x = as_variable(x)

    if dezero.core.Config.train:
        mask = xp.random.rand(*x.shape) > dropout_ratio
        scale = xp.array(1.0 - dropout_ratio).astype(x.dtype)
        y = x * mask / scale
        return y
    else:
        #测试模式
        return x